from ultralytics import YOLO
from swanlab.integration.ultralytics import add_swanlab_callback


if __name__ == "__main__":
    model_name = "yolo11s"
    model = YOLO(model_name + ".yaml")
    model.load()
    # 添加swanlab回调
    add_swanlab_callback(
        model,
        project="VID60"
    )

    model.train(
        data="/home/amov/spirecv-training-pro/ultralytics/ultralytics/cfg/datasets/mav.yaml",
        epochs=100,
        imgsz=640,
        batch=16,
        pretrained=model_name + '.pt',
        copy_paste=0,
        mosaic=0,
        workers=4
    )
